"""
按比例拆分文件
"""
import os
import shutil
import random

if __name__ == '__main__':
    root_dir = "/media/hsmy/wanghao_18T/20240607/yolo/all_/"
    output_dir = "/media/hsmy/wanghao_18T/20240607/yolo/all/"
    train_ratio = 0.8

    png_files = os.listdir(root_dir + "images")
    random.shuffle(png_files)

    # Compute the split indices
    total_files = len(png_files)
    train_end = int(total_files * train_ratio)

    # Split the files into train, val, and test
    train_files = png_files[:train_end]
    val_files = png_files[train_end:]


    def move_files(files, dest):
        for file in files:
            file_name = file.split('.')[0]
            png_name = f"{file_name}.png"
            txt_name = f"{file_name}.txt"

            png_src_path = os.path.join(root_dir, "images", png_name)
            label_src_path = os.path.join(root_dir, "labels", txt_name)

            png_dest_path = os.path.join(output_dir, "images", dest, png_name)
            label_dest_path = os.path.join(output_dir, "labels", dest, txt_name)

            shutil.move(png_src_path, png_dest_path)
            if os.path.exists(label_src_path):
                shutil.move(label_src_path, label_dest_path)

    os.makedirs(output_dir + "/images/train", exist_ok=True)
    os.makedirs(output_dir + "/images/val", exist_ok=True)

    os.makedirs(output_dir + "/labels/train", exist_ok=True)
    os.makedirs(output_dir + "/labels/val", exist_ok=True)

    # Move the files
    move_files(train_files, "train")
    move_files(val_files, "val")
